import os
from os.path import join, dirname
import numpy as np
import trimesh


def create_and_save_urdf(obj_bbox, obj_aug_strategy, urdf_save_path):

    obj_lengths = obj_bbox[1] - obj_bbox[0]
    xyz = obj_aug_strategy / 2
    xyz[0], xyz[2] = 0, 0
    size = obj_lengths + obj_aug_strategy

    os.makedirs(dirname(urdf_save_path), exist_ok=True)
    with open(urdf_save_path, "w") as f:
        f.write("<?xml version=\"1.0\"?>\n")
        f.write("<robot name=\"object\">\n")
        f.write("\t<link name=\"object\">\n")
        f.write("\t\t<visual>\n")
        f.write("\t\t\t<origin\n")
        f.write("\t\t\t\txyz=\"{}\"\n".format(" ".join([str(x) for x in list(xyz)])))
        f.write("\t\t\t\trpy=\"0 0 0\" />\n")
        f.write("\t\t\t<geometry>\n")
        f.write("\t\t\t\t<box size=\"{}\" />\n".format(" ".join([str(x) for x in list(size)])))
        f.write("\t\t\t</geometry>\n")
        f.write("\t\t</visual>\n")
        f.write("\t\t<collision>\n")
        f.write("\t\t\t<origin\n")
        f.write("\t\t\t\txyz=\"{}\"\n".format(" ".join([str(x) for x in list(xyz)])))
        f.write("\t\t\t\trpy=\"0 0 0\" />\n")
        f.write("\t\t\t<geometry>\n")
        f.write("\t\t\t\t<box size=\"{}\" />\n".format(" ".join([str(x) for x in list(size)])))
        f.write("\t\t\t</geometry>\n")
        f.write("\t\t</collision>\n")
        f.write("\t\t<inertial>\n")
        f.write("\t\t\t<mass value=\"0.5\"/>\n")
        f.write("\t\t\t<inertia ixx=\"0.05\" ixy=\"0.0\" ixz=\"0.0\" iyy=\"0.05\" iyz=\"0.0\" izz=\"0.05\"/>\n")
        f.write("\t\t</inertial>\n")
        f.write("\t</link>\n")
        f.write("</robot>\n")


def create_and_save_mesh(obj_bbox, obj_aug_strategy, mesh_save_path):

    obj_lengths = obj_bbox[1] - obj_bbox[0]
    xyz = obj_aug_strategy / 2
    xyz[0], xyz[2] = 0, 0
    size = obj_lengths + obj_aug_strategy

    os.makedirs(dirname(mesh_save_path), exist_ok=True)
    T = np.eye(4)
    T[:3, 3] = xyz
    bbox = trimesh.primitives.Box(extents=size, transform=T)
    bbox.export(mesh_save_path)
